# 导入库
from joblib import load
import numpy as np

class ClassificationModel(object):
    def __init__(self):
        self.knn = load('knn.jobilb')
        self.scaler = load('scaler.jobilb')
        self.label_encoder = load('label_encoder.jobilb')

    def get_predictor_category(self, new_data):
        """
        获取预测的分类
        """
        # 转换为数组并保存特征顺序
        feature_order = ['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 'fcff', 'fcfe',
                         'tangible_asset', 'bps', 'grossprofit_margin', 'npta']
        new_values = np.array([new_data[feature] for feature in feature_order])
        # 标准化新数据
        new_values_scaled = self.scaler.transform(new_values.reshape(1, -1))
        # 预测分类
        predicted_label = self.knn.predict(new_values_scaled)
        # 转换为原始分类标签
        predicted_category = self.label_encoder.inverse_transform(predicted_label)
        return predicted_category[0]

if __name__ == '__main__':
    ci = ClassificationModel()
    new_data = {
        'eps': -4.17,
        'total_revenue_ps': 28.7641,
        'undist_profit_ps': 3.6095,
        'gross_margin': 34911600000,
        'fcff': 134828000000,
        'fcfe': 133615000000,
        'tangible_asset': 136797000000,
        'bps': 16.987,
        'grossprofit_margin': 10.1731,
        'npta': -3.4899
    }
    predicted_category = ci.get_predictor_category(new_data)
    print(f"预测的分类: {predicted_category}")